#!/usr/bin/env python3
# Copyright (c) 2019-2020 Windy Albert

import base58check
import ecdsa
import os
import json
import sys
import hashlib


COLOR_NONE  = "\033[0m"
COLOR_RED   = "\033[0;31m"
COLOR_GREEN = "\033[0;32m"
COLOR_YELLOW = "\033[0;33m"

def usage():
    print("Make by Windy Albert")
    print("Usage:")
    print("\t./sign [keyfile] [datafile]")
    
def main():
    if len(sys.argv) != 3 :
        usage()
        return

    kfile_name = sys.argv[1]
    dfile_name = sys.argv[2]

    if (not os.path.exists(kfile_name)) or (not os.path.exists(dfile_name)):
        print("There is no prikey or datafile exist, please re-choose correctlly !!")
        return

    # Get prikey ...
    f = open(kfile_name,'r')
    d = f.read()
    wif = json.loads(d)
    f.close()

    wif_hex = base58check.b58decode(wif["wif"])
    sk = ecdsa.SigningKey.from_string( bytes.fromhex(wif_hex[1:33].hex()), curve=ecdsa.SECP256k1)    

    # Get data to sign ...
    f = open(dfile_name,'r')
    d = f.read()
    data_raw = json.loads(d)
    f.close()

    sha256  = hashlib.new('sha256',bytes.fromhex(data_raw["data"]))
    if data_raw["sha256"] != sha256.hexdigest():
        print("Error : hash is wrong ...quit ! ")
        return 

    s = sk.sign(data=bytes.fromhex(data_raw["data"]),hashfunc=hashlib.sha256)
    data_raw["sign"] = s.hex()

    f = open(dfile_name,'w')
    f.write(json.dumps(data_raw))
    f.close()
    print("Sign "+COLOR_GREEN + dfile_name+COLOR_NONE+" with " + COLOR_GREEN +kfile_name+COLOR_NONE+ " successfully !")



if __name__ == '__main__':
    main()